package main

import (
	"crypto/sha256"
	"fmt"
)

// 计算hash
func calHash(data []byte) []byte {
	hash := sha256.Sum256(data)
	return hash[:]
}

// 创建merkleTree节点
type MerkleNode struct {
	hash       []byte
	leftChild  *MerkleNode
	rightChild *MerkleNode
}

// 创建叶子节点
func NewMerkleLeaf(data []byte) *MerkleNode {
	return &MerkleNode{
		hash: calHash(data),
	}
}

// 创建父节点
func NewMerkleParent(left *MerkleNode, right *MerkleNode) *MerkleNode {
	parent := &MerkleNode{
		leftChild:  left,
		rightChild: right,
	}

	// 计算父节点的hash值
	parent.hash = calHash(append(left.hash, right.hash...))
	return parent
}

// 创建Merkle树
func BuildMerkleTree(data [][]byte) *MerkleNode {
	var nodes []*MerkleNode

	// 创建叶子节点
	for _, d := range data {
		nodes = append(nodes, NewMerkleLeaf(d))
	}

	// 如果叶子节点数量为奇数，复制最后一个节点以保证数量为偶数
	if len(nodes)%2 != 0 {
		nodes = append(nodes, nodes[len(nodes)-1])
	}

	// 构建Merkle树
	for len(nodes) > 1 {
		var newLevel []*MerkleNode

		for i := 0; i < len(nodes); i += 2 {
			newLevel = append(newLevel, NewMerkleParent(nodes[i], nodes[i+1]))
		}

		nodes = newLevel
	}

	return nodes[0]
}

func main() {
	// 示例数据
	data := [][]byte{
		[]byte("data1"),
		[]byte("data2"),
		[]byte("data3"),
		[]byte("data4"),
	}

	// 构建Merkle树
	root := BuildMerkleTree(data)

	// 打印根节点的hash值
	fmt.Printf("Merkle Root Hash: %x\n", root.hash)
}
